import type { FullConfig, Reporter, Suite, TestCase } from './packages/playwright-test/reporter';

export default class ShardWeightReporter implements Reporter {
  suite!: Suite;
  config!: FullConfig;

  onBegin(config: FullConfig, suite: Suite): void {
    this.suite = suite;
    this.config = config;
  }

  onEnd() {
    const tests = this.suite.allTests();

    const testsByBot = Object.entries(this.groupTestsByBot(tests));
    testsByBot.sort((a, b) => a[0].localeCompare(b[0]));

    console.log('Suggested shard weights per bot:');
    for (const [botName, botTests] of testsByBot) {
      const suggestedWeights = this.calculateShardWeights(botTests);
      console.log(`${botName}:`, this.normaliseWeights(suggestedWeights));
    }
  }

  private groupTestsByBot(tests: TestCase[]): Record<string, TestCase[]> {
    const testsByBot: Record<string, TestCase[]> = {};

    for (const test of tests) {
      let botName = test.tags[0] ?? 'unknown';
      botName = botName.replace(/-\d+$/, ''); // TODO: fix on playwright end. bot name shouldn't contain shard index suffix
      testsByBot[botName] ??= [];
      testsByBot[botName].push(test);
    }

    return testsByBot;
  }

  private sortByTestGroups(tests: TestCase[]) {
    // recreates the sorting that filterForShard gets as input
    tests.sort((a, b) => {
      const shardIndexA = a.results[0]!.shardIndex;
      const shardIndexB = b.results[0]!.shardIndex;
      if (shardIndexA !== shardIndexB)
        return shardIndexA - shardIndexB;

      const startTimeA = a.results[0]!.startTime.getTime();
      const startTimeB = b.results[0]!.startTime.getTime();
      return startTimeA - startTimeB;
    });

  }

  private calculateShardWeights(tests: TestCase[]): number[] {
    if (tests.length === 0)
      return [];

    this.sortByTestGroups(tests);

    const shardTotal = tests[tests.length - 1].results[0]!.shardIndex;

    const totalDuration = this.sum(tests.map(test => test.results[0]!.duration));
    const optimalDuration = totalDuration / shardTotal;

    const weights: number[] = [];
    let currentDuration = 0;
    let currentTestCount = 0;

    for (const test of tests) {
      const duration = test.results[0]!.duration;
      currentDuration += duration;
      currentTestCount++;

      // When we've reached the optimal shard duration, record this shard's weight
      if (currentDuration >= optimalDuration && weights.length < shardTotal - 1) {
        weights.push(currentTestCount);
        currentDuration = 0;
        currentTestCount = 0;
      }
    }
    weights.push(currentTestCount);

    return weights;
  }

  private normaliseWeights(weights: number[]): number[] {
    const total = weights.reduce((a, b) => a + b, 0);
    weights = weights.map(w => Math.floor((w / total) * 100));
    const remaining = 100 - weights.reduce((a, b) => a + b, 0);
    for (let i = 0; i < remaining; i++)
      weights[i % weights.length]++;
    return weights;
  }

  private sum(array: number[]): number {
    return array.reduce((a, b) => a + b, 0);
  }
}